/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         https://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef itkObjectToObjectMetric_hxx
#define itkObjectToObjectMetric_hxx

#include "itkTransform.h"
#include "itkIdentityTransform.h"
#include "itkCompositeTransform.h"

namespace itk
{

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::ObjectToObjectMetric()

{
  // Both transforms default to an identity transform.
  using MovingIdentityTransformType = IdentityTransform<TParametersValueType, Self::MovingDimension>;
  using FixedIdentityTransformType = IdentityTransform<TParametersValueType, Self::FixedDimension>;
  this->m_FixedTransform = FixedIdentityTransformType::New();
  this->m_MovingTransform = MovingIdentityTransformType::New();

  this->m_VirtualImage = nullptr;

  this->m_UserHasSetVirtualDomain = false;
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::Initialize()
{
  if (!this->m_FixedTransform)
  {
    itkExceptionMacro("Fixed transform is not present");
  }

  if (!this->m_MovingTransform)
  {
    itkExceptionMacro("Moving transform is not present");
  }

  /* Special checks for when the moving transform is dense/high-dimensional */
  if (this->HasLocalSupport())
  {
    /* Verify that virtual domain and displacement field are the same size
     * and in the same physical space. Handles CompositeTransform by checking
     * if first applied transform is DisplacementFieldTransform */
    this->VerifyDisplacementFieldSizeAndPhysicalSpace();

    /* Verify virtual image pixel type is scalar. This effects the calculation
     * of offsets in ComputeParameterOffsetFromVirtualIndex().
     * NOTE:  Can this be checked at compile time? ConceptChecking has a
     * HasPixelTraits class, but looks like it just verifies that type T
     * has PixelTraits associated with it, and not a particular value.
     */
    if (PixelTraits<VirtualPixelType>::Dimension != 1)
    {
      itkExceptionMacro("VirtualPixelType must be scalar for use "
                        "with high-dimensional transform. "
                        "Dimensionality is "
                        << PixelTraits<VirtualPixelType>::Dimension);
    }
  }
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::SetTransform(
  MovingTransformType * transform)
{
  this->SetMovingTransform(transform);
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
const typename ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  MovingTransformType *
  ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::GetTransform()
{
  return this->GetMovingTransform();
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::UpdateTransformParameters(
  const DerivativeType & derivative,
  TParametersValueType   factor)
{
  // Rely on transform::UpdateTransformParameters to verify proper
  // size of derivative.
  this->m_MovingTransform->UpdateTransformParameters(derivative, factor);
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
typename ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  NumberOfParametersType
  ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::GetNumberOfParameters()
    const
{
  return this->m_MovingTransform->GetNumberOfParameters();
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
const typename ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  ParametersType &
  ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::GetParameters() const
{
  return this->m_MovingTransform->GetParameters();
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::SetParameters(
  ParametersType & params)
{
  this->m_MovingTransform->SetParametersByValue(params);
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
typename ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  NumberOfParametersType
  ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
    GetNumberOfLocalParameters() const
{
  return this->m_MovingTransform->GetNumberOfLocalParameters();
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
bool
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::HasLocalSupport() const
{
  return (this->m_MovingTransform->GetTransformCategory() ==
          MovingTransformType::TransformCategoryEnum::DisplacementField);
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
bool
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  TransformPhysicalPointToVirtualIndex(const VirtualPointType & point, VirtualIndexType & index) const
{
  if (this->m_VirtualImage)
  {
    return this->m_VirtualImage->TransformPhysicalPointToIndex(point, index);
  }

  itkExceptionMacro("m_VirtualImage is undefined. Cannot transform.");
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  TransformVirtualIndexToPhysicalPoint(const VirtualIndexType & index, VirtualPointType & point) const
{
  if (this->m_VirtualImage)
  {
    this->m_VirtualImage->TransformIndexToPhysicalPoint(index, point);
  }
  else
  {
    itkExceptionMacro("m_VirtualImage is undefined. Cannot transform.");
  }
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::SetVirtualDomain(
  const VirtualSpacingType &   spacing,
  const VirtualOriginType &    origin,
  const VirtualDirectionType & direction,
  const VirtualRegionType &    region)
{
  if (this->m_VirtualImage.IsNull() || (this->m_VirtualImage->GetSpacing() != spacing) ||
      (this->m_VirtualImage->GetOrigin() != origin) || (this->m_VirtualImage->GetDirection() != direction) ||
      (this->m_VirtualImage->GetLargestPossibleRegion() != region) ||
      (this->m_VirtualImage->GetBufferedRegion() != region))
  {
    this->m_VirtualImage = VirtualImageType::New();
    this->m_VirtualImage->SetSpacing(spacing);
    this->m_VirtualImage->SetOrigin(origin);
    this->m_VirtualImage->SetDirection(direction);
    this->m_VirtualImage->SetRegions(region);
    this->m_UserHasSetVirtualDomain = true;
    this->Modified();
  }
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::SetVirtualDomainFromImage(
  const VirtualImageType * virtualImage)
{
  this->SetVirtualDomain(virtualImage->GetSpacing(),
                         virtualImage->GetOrigin(),
                         virtualImage->GetDirection(),
                         virtualImage->GetLargestPossibleRegion());
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
const TimeStamp &
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  GetVirtualDomainTimeStamp() const
{
  if (!this->GetVirtualImage())
  {
    return this->GetTimeStamp();
  }

  if (this->GetTimeStamp() > this->GetVirtualImage()->GetTimeStamp())
  {
    return this->GetTimeStamp();
  }

  return this->GetVirtualImage()->GetTimeStamp();
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
bool
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::IsInsideVirtualDomain(
  const VirtualPointType & point) const
{
  if (!this->m_VirtualImage.IsNull())
  {
    const auto index = m_VirtualImage->TransformPhysicalPointToIndex(point);
    return this->GetVirtualRegion().IsInside(index);
  }

  // Otherwise always return true since a virtual domain hasn't been defined, and
  // we assume the user is working in an unconstrained domain.
  return true;
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
bool
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::IsInsideVirtualDomain(
  const VirtualIndexType & index) const
{
  if (!this->m_VirtualImage.IsNull())
  {
    return this->GetVirtualRegion().IsInside(index);
  }

  // Otherwise always return true since a virtual domain hasn't been defined, and
  // we assume the user is working in an unconstrained domain.
  return true;
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
OffsetValueType
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  ComputeParameterOffsetFromVirtualPoint(const VirtualPointType &       point,
                                         const NumberOfParametersType & numberOfLocalParameters) const
{
  if (!this->m_VirtualImage.IsNull())
  {
    VirtualIndexType index;
    if (!this->m_VirtualImage->TransformPhysicalPointToIndex(point, index))
    {
      itkExceptionMacro(" point is not inside virtual domain. Cannot compute offset. ");
    }
    return this->ComputeParameterOffsetFromVirtualIndex(index, numberOfLocalParameters);
  }

  itkExceptionMacro("m_VirtualImage is undefined. Cannot calculate offset.");
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
OffsetValueType
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  ComputeParameterOffsetFromVirtualIndex(const VirtualIndexType &       index,
                                         const NumberOfParametersType & numberOfLocalParameters) const
{
  if (m_VirtualImage)
  {
    const OffsetValueType offset = this->m_VirtualImage->ComputeOffset(index) * numberOfLocalParameters;
    return offset;
  }

  itkExceptionMacro("m_VirtualImage is undefined. Cannot calculate offset.");
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
typename ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  VirtualSpacingType
  ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::GetVirtualSpacing()
    const
{
  if (this->m_VirtualImage)
  {
    return this->m_VirtualImage->GetSpacing();
  }

  auto spacing = MakeFilled<VirtualSpacingType>(NumericTraits<typename VirtualSpacingType::ValueType>::OneValue());
  return spacing;
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
typename ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  VirtualDirectionType
  ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::GetVirtualDirection()
    const
{
  if (this->m_VirtualImage)
  {
    return this->m_VirtualImage->GetDirection();
  }

  auto direction =
    MakeFilled<VirtualDirectionType>(NumericTraits<typename VirtualDirectionType::ValueType>::OneValue());
  return direction;
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
auto
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::GetVirtualOrigin() const
  -> VirtualOriginType
{
  if (this->m_VirtualImage)
  {
    return this->m_VirtualImage->GetOrigin();
  }

  VirtualOriginType origin{};
  return origin;
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
const typename ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  VirtualRegionType &
  ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::GetVirtualRegion() const
{
  if (this->m_VirtualImage)
  {
    return this->m_VirtualImage->GetBufferedRegion();
  }

  itkExceptionMacro("m_VirtualImage is undefined. Cannot return region. ");
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
const typename ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  MovingDisplacementFieldTransformType *
  ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
    GetMovingDisplacementFieldTransform() const
{
  // If it's a composite transform and the displacement field is the first
  // to be applied (i.e. the most recently added), then return that.
  using MovingCompositeTransformType = CompositeTransform<CoordinateRepresentationType, Self::MovingDimension>;
  const MovingTransformType * transform = this->m_MovingTransform.GetPointer();
  // If it's a CompositeTransform, get the last transform (1st applied).
  const auto * comptx = dynamic_cast<const MovingCompositeTransformType *>(transform);
  if (comptx != nullptr)
  {
    transform = comptx->GetBackTransform();
  }
  // Cast to a DisplacementField type.
  const auto * deftx = dynamic_cast<const MovingDisplacementFieldTransformType *>(transform);
  return deftx;
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::
  VerifyDisplacementFieldSizeAndPhysicalSpace()
{
  // TODO: replace with a common external method to check this,
  // possibly something in Transform.

  /* Verify that virtual domain and displacement field are the same size
   * and in the same physical space.
   * Effects transformation, and calculation of offset in StoreDerivativeResult.
   * If it's a composite transform and the displacement field is the first
   * to be applied (i.e. the most recently added), then it has to be
   * of the same size, otherwise not.
   * Eventually we'll want a method in Transform something like a
   * GetInputDomainSize to check this cleanly. */
  const MovingDisplacementFieldTransformType * displacementTransform = this->GetMovingDisplacementFieldTransform();
  if (displacementTransform == nullptr)
  {
    itkExceptionMacro("Expected the moving transform to be of type DisplacementFieldTransform or derived, "
                      "or a CompositeTransform with DisplacementFieldTransform as the last to have been added.");
  }
  using FieldType = typename MovingDisplacementFieldTransformType::DisplacementFieldType;
  const typename FieldType::ConstPointer field = displacementTransform->GetDisplacementField();
  const typename FieldType::RegionType   fieldRegion = field->GetBufferedRegion();
  const VirtualRegionType                virtualRegion = this->GetVirtualRegion();
  if (virtualRegion.GetSize() != fieldRegion.GetSize() || virtualRegion.GetIndex() != fieldRegion.GetIndex())
  {
    itkExceptionMacro(
      "Virtual domain and moving transform displacement field"
      " must have the same size and index for BufferedRegion."
      << std::endl
      << "Virtual size/index: " << virtualRegion.GetSize() << " / " << virtualRegion.GetIndex() << std::endl
      << "Displacement field size/index: " << fieldRegion.GetSize() << " / " << fieldRegion.GetIndex() << std::endl);
  }

  /* check that the image occupy the same physical space, and that
   * each index is at the same physical location.
   * this code is from ImageToImageFilter */

  /* tolerance for origin and spacing depends on the size of pixel
   * tolerance for directions a fraction of the unit cube. */
  const double     coordinateTol = 1.0e-6 * this->GetVirtualSpacing()[0];
  constexpr double directionTol = 1.0e-6;

  if (!this->GetVirtualOrigin().GetVnlVector().is_equal(field->GetOrigin().GetVnlVector(), coordinateTol) ||
      !this->GetVirtualSpacing().GetVnlVector().is_equal(field->GetSpacing().GetVnlVector(), coordinateTol) ||
      !this->GetVirtualDirection().GetVnlMatrix().is_equal(field->GetDirection().GetVnlMatrix(), directionTol))
  {
    std::ostringstream originString;

    originString << "Virtual Origin: " << this->GetVirtualOrigin()
                 << ", DisplacementField Origin: " << field->GetOrigin() << std::endl;
    std::ostringstream spacingString;
    spacingString << "Virtual Spacing: " << this->GetVirtualSpacing()
                  << ", DisplacementField Spacing: " << field->GetSpacing() << std::endl;
    std::ostringstream directionString;
    directionString << "Virtual Direction: " << this->GetVirtualDirection()
                    << ", DisplacementField Direction: " << field->GetDirection() << std::endl;
    itkExceptionMacro("Virtual Domain and DisplacementField do not "
                      << "occupy the same physical space! You may be able to "
                      << "simply call displacementField->CopyInformation( "
                      << "metric->GetVirtualImage() ) to align them. " << std::endl
                      << originString.str() << spacingString.str() << directionString.str());
  }
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
bool
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::VerifyNumberOfValidPoints(
  MeasureType &    value,
  DerivativeType & derivative) const
{
  if (this->m_NumberOfValidPoints == 0)
  {
    value = NumericTraits<MeasureType>::max();
    derivative.Fill(DerivativeValueType{});
    itkWarningMacro("No valid points were found during metric evaluation. "
                    "For image metrics, verify that the images overlap appropriately. "
                    "For instance, you can align the image centers by translation. "
                    "For point-set metrics, verify that the fixed points, once transformed "
                    "into the virtual domain space, actually lie within the virtual domain.");
    return false;
  }
  return true;
}

template <unsigned int TFixedDimension,
          unsigned int TMovingDimension,
          typename TVirtualImage,
          typename TParametersValueType>
void
ObjectToObjectMetric<TFixedDimension, TMovingDimension, TVirtualImage, TParametersValueType>::PrintSelf(
  std::ostream & os,
  Indent         indent) const
{
  Superclass::PrintSelf(os, indent);

  itkPrintSelfObjectMacro(FixedTransform);
  itkPrintSelfObjectMacro(MovingTransform);
  itkPrintSelfObjectMacro(VirtualImage);

  itkPrintSelfBooleanMacro(UserHasSetVirtualDomain);
  os << indent
     << "NumberOfValidPoints: " << static_cast<typename NumericTraits<SizeValueType>::PrintType>(m_NumberOfValidPoints)
     << std::endl;
}

} // namespace itk

#endif
